package com.test.demo.config;

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * 动态数据源
 */
@Component
public class DynamicDataSource extends AbstractRoutingDataSource {

    public DynamicDataSource(DataSource defaultTargetDataSource) {
        super.setDefaultTargetDataSource(defaultTargetDataSource);
        // 将Map<String, DataSource>转换为Map<Object, Object>
        Map<Object, Object> targetDataSources = new HashMap<>();
        DynamicDataSourceContextHolder.getDataSourceMap().forEach((key, dataSource) -> 
            targetDataSources.put(key, dataSource)
        );
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getDataSourceKey();
    }

    /**
     * 添加数据源
     */
    public void addDataSource(String key, DataSource dataSource) {
        DynamicDataSourceContextHolder.addDataSource(key, dataSource);
        // 重新加载数据源
        Map<Object, Object> targetDataSources = new HashMap<>();
        DynamicDataSourceContextHolder.getDataSourceMap().forEach((k, v) -> 
            targetDataSources.put(k, v)
        );
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }

    /**
     * 删除数据源
     */
    public void removeDataSource(String key) {
        DynamicDataSourceContextHolder.closeDataSource(key);
        // 重新加载数据源
        Map<Object, Object> targetDataSources = new HashMap<>();
        DynamicDataSourceContextHolder.getDataSourceMap().forEach((k, v) -> 
            targetDataSources.put(k, v)
        );
        super.setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();
    }
} 